import math
import yaml
import numpy as np
from utils import set_and_load_parameters_imp_2


def padding(data, length: int, value:int):
    """ 长度补齐 """
    if isinstance(data, np.ndarray):
        shape = list(data.shape)
        shape[0] = length - shape[0]
        paddings = value * np.ones(shape, dtype=np.float32)
        data = np.concatenate([data, paddings])
    elif isinstance(data, list):
        paddings = [value] * (length - len(data))
        data.extend(paddings)
    return data


def load_yaml(yaml_path: str = 'configs.yaml'):
    """ 读配置文件, 默认configs.yaml """
    with open(yaml_path, 'r',encoding = 'utf-8') as f:
        data = yaml.safe_load(f)
    return data


class PairWithSpanMaskingScheme:
    def __init__(self, params=None):
        self.params = params
        """边界预测的最大长度"""
        self.max_pair_targets = self.params['max_pair_targets']
        self.lower = self.params['span_lower']
        self.upper = self.params['span_upper']
        self.mask_ratio = self.params['mask_factor']
        self.pad = 0
        self.mask_id = -9999
        self.boundpoints = 2
        self.return_only_spans = False
        self.lens = list(range(self.lower, self.upper + 1))
        self.p = self.params['geometric_p']
        self.len_distrib = [self.p * (1-self.p)**(i - self.lower) for i in range(self.lower, self.upper + 1)] if self.p >= 0 else None
        self.len_distrib = [x / (sum(self.len_distrib)) for x in self.len_distrib]
        # print(self.len_distrib, self.lens)


    def mask(self, sentence):
        """mask tokens for masked language model training
        Args:
            sentence: 1d tensor, token list to be masked
            mask_ratio: ratio of tokens to be masked in the sentence
        Return:
            masked_sent: masked sentence
        """
        sent_length = len(sentence)
        mask_num = math.ceil(sent_length * self.mask_ratio)
        mask = set()
        mask_indicator = set()
        spans = []
        while len(mask) < mask_num:
            span_len = np.random.choice(self.lens, p=self.len_distrib)
            anchor = np.random.choice(range(self.boundpoints, sent_length-self.boundpoints-1))
            if anchor in mask_indicator:
                continue
            # find word start, end
            left1 = anchor
            right1 = left1 + span_len
            if right1 > sent_length-self.boundpoints:
                continue
            spans.append([left1, left1])
            for i in range(left1, right1):
                if len(mask) >= mask_num:
                    spans, mask, mask_indicator = merge_intervals(spans, self.max_pair_targets)
                    break
                mask.add(i)
                mask_indicator.add(i)
                spans[-1][-1] = i
            spans, mask, mask_indicator = merge_intervals(spans, self.max_pair_targets)
        sentence_target, masked_sentence, pair_targets = span_masking(sentence, spans, self.pad, self.mask_id, self.max_pair_targets, mask, self.boundpoints)
        mask_list = list(mask)
        masked_label = np.ones(sentence.shape[0])
        masked_label[mask_list] = 0
        if self.return_only_spans:
            pair_targets = None
        return sentence_target, masked_sentence, pair_targets, masked_label


def span_masking(sentence, spans, pad, mask_id, pad_len, mask, boundpoints=1):
    """endpoints设置为'external'，span的端点会被包含在内"""
    sentence_copy = sentence.copy()
    sent_length = len(sentence_copy)
    # target = [pad] * sent_length
    pair_targets = []
    # spans, mask_set = merge_intervals(spans, pad_len)
    assert len(mask) == sum([e - s + 1 for s, e in spans])
    for start, end in spans:
        lower_limit = max(0, start - boundpoints)
        upper_limit = min(sent_length - 1, end + boundpoints)
        if start > lower_limit and end < upper_limit:
            stacked_array = np.stack([sentence_copy[i] for i in range(start, end + 1)])
            padded_array = padding(stacked_array, pad_len, pad)
            pair_targets.append(
                [i for i in range(start - boundpoints, start)] +
                [i for i in range(end + 1, end + boundpoints + 1)] +
                [padded_array])

        for i in range(start, end + 1):
            assert i in mask
            rand = np.random.random()
            if rand < 0.8:
                sentence_copy[i] = mask_id

    return sentence, sentence_copy, pair_targets


def merge_intervals(intervals, max_pair_targets):
    intervals = sorted(intervals, key=lambda x : x[0])
    merged = []
    mask = set()  # 初始化一个空的mask
    mask_indicator = set()
    for interval in intervals:
        if not merged or merged[-1][1] + 2 < interval[0]:
            merged.append(list(interval))  # convert tuple to list
        else:
            merged[-1][1] = max(merged[-1][1], interval[1])

        # Check if the merged interval is longer than max_pair_targets
        if merged[-1][1] - merged[-1][0] + 1 > max_pair_targets:
            # If so, truncate it to max_pair_targets
            merged[-1][1] = merged[-1][0] + max_pair_targets - 1

        # Add the merged interval to mask
        mask.update(range(merged[-1][0], merged[-1][1] + 1))
        mask_indicator.update(range(merged[-1][0], merged[-1][1] + 3))

    return merged, mask, mask_indicator


def pad_to_len(pair_targets, pad, max_pair_target_len):
    for i in range(len(pair_targets)):
        pair_targets[i] = pair_targets[i][:max_pair_target_len]
        this_len = len(pair_targets[i])
        for j in range(max_pair_target_len - this_len - 1):
            pair_targets[i].append(pad)
    return pair_targets


if __name__ == "__main__":
    cfgs = load_yaml()
    sentence = np.random.rand(100, 5)

    scheme = PairWithSpanMaskingScheme(cfgs)
    target_sentence, masked_sentence, pair_targets, mask = scheme.mask(sentence)
    """
    print("target_sentence:", type(target_sentence), target_sentence.shape)
    print(target_sentence)
    print("masked_sentence:", type(masked_sentence), masked_sentence.shape)
    print(masked_sentence)
    """
    # print("Pair targets:", type(pair_targets))
    # print(pair_targets)

    pairs = [item[:4] for item in pair_targets]
    print('pairs', pairs)
    # pair_targets = [item[4] for item in pair_targets]
    # pair_targets = torch.Tensor(pair_targets)
    # print('pair_targets', type(pair_targets))
    # print(pair_targets)
    print('mask', len(mask), mask)
    """
    mask_list = list(mask)
    label = np.ones(sentence.shape[0])
    label[mask_list] = 0
    label = torch.LongTensor(label).unsqueeze(0)
    print('label', label.shape)
    print(label)
    """